#include <hip/hip_runtime.h>
#include <iostream>
#include <mutex>
#include <vector>
#include <thread>
#include <unistd.h>
#include <sys/wait.h>

// 全局互斥锁用于控制输出
std::mutex g_console_mutex;

#define CHECK(cmd) \
{ \
    hipError_t error = cmd; \
    if (error != hipSuccess) { \
        std::cerr << "Error: " << hipGetErrorString(error) \
                  << " (" << __FILE__ << ":" << __LINE__ << ")" << std::endl; \
        exit(EXIT_FAILURE); \
    } \
}

__global__ void vectorAdd(const float* A, const float* B, float* C, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        C[i] = A[i] + B[i];
    }
}

// 通过hy-smi命令获取设备数量
int getDeviceCountFromHySmi() {
    std::array<char, 128> buffer;
    std::string output;
    FILE* pipe = popen("hy-smi -i", "r");
    if (!pipe) {
        throw std::runtime_error("popen() failed!");
    }
    while (fgets(buffer.data(), buffer.size(), pipe) != nullptr) {
        output += buffer.data();
    }
    pclose(pipe);

    int count = 0;
    std::istringstream iss(output);
    std::string line;
    while (std::getline(iss, line)) {
        if (line.find("Device ID") != std::string::npos) {
            ++count;
        }
    }
    return count;
}

// 解析环境变量HIP_VISIBLE_DEVICES中的设备编号列表，计算 GPU 卡个数
int getDeviceCountFromEnv() {
    const char* envVar = std::getenv("HIP_VISIBLE_DEVICES");
    if (envVar == nullptr) {  // 环境变量未设置则通过hy-smi获取
        return getDeviceCountFromHySmi();
    }

    std::string deviceList(envVar);
    std::stringstream ss(deviceList);
    std::string token;
    int count = 0;

    // 按逗号分割设备编号列表
    while (std::getline(ss, token, ',')) {
        if (!token.empty()) {
            ++count;
        }
    }

    return count;
}

void gpuWorker(int device_id) {
    // 设置当前线程使用的设备
    CHECK(hipSetDevice(device_id));

    hipDeviceProp_t prop;
    CHECK(hipGetDeviceProperties(&prop, device_id));
    {
        std::lock_guard<std::mutex> lock(g_console_mutex);
        std::cout << "Process " << getpid()
                  << " Thread " << std::this_thread::get_id()
                  << " using GPU " << device_id
                  << " (" << prop.name << ")" << std::endl;
    }

    constexpr int N = 64 * 1024 * 1024;
    constexpr size_t bytes = N * sizeof(float);

    // 分配内存
    float *A, *B, *C;
    CHECK(hipMalloc(&A, bytes));
    CHECK(hipMalloc(&B, bytes));
    CHECK(hipMalloc(&C, bytes));

    // 初始化数据
    std::vector<float> h_A(N, 1.0f), h_B(N, 2.0f), h_C(N, 0.0f);

    CHECK(hipMemcpy(A, h_A.data(), bytes, hipMemcpyHostToDevice));
    CHECK(hipMemcpy(B, h_B.data(), bytes, hipMemcpyHostToDevice));

    // 执行核函数
    constexpr int blockSize = 256;
    constexpr int gridSize = (N + blockSize - 1) / blockSize;

    hipLaunchKernelGGL(vectorAdd,
                      dim3(gridSize), dim3(blockSize), 0, 0,
                      A, B, C, N);

    // 验证结果
    CHECK(hipMemcpy(h_C.data(), C, bytes, hipMemcpyDeviceToHost));

    for (int i = 0; i < N; ++i) {
        if (fabs(h_C[i] - 3.0f) > 1e-6) {
            std::cerr << "Verification failed on device " << device_id << std::endl;
            exit(EXIT_FAILURE);
        }
    }

    // 释放资源
    CHECK(hipFree(A));
    CHECK(hipFree(B));
    CHECK(hipFree(C));
    {
        std::lock_guard<std::mutex> lock(g_console_mutex);
        std::cout << "Finishing gpuWorker on device " << device_id
                  << " in process " << getpid()
                  << " and thread " << std::this_thread::get_id() << std::endl;
    }
}

int main() {
    int device_count = 0;
    // CHECK(hipGetDeviceCount(&device_count));
    device_count = getDeviceCountFromEnv();
    if (device_count < 1) {
        std::cerr << "No HIP devices found" << std::endl;
        return EXIT_FAILURE;
    }

    std::cout << "Found " << device_count << " GPUs\n";

    // 创建子进程执行用例
    std::vector<pid_t> pids;
    for (int dev = 0; dev < device_count-1; ++dev) {
        pid_t pid = fork();
        if (pid == 0) { // 子进程
            // CHECK(hipSetDevice(dev));
            gpuWorker(dev);
            exit(EXIT_SUCCESS);
        } else if (pid > 0) { // 父进程
            pids.push_back(pid);
        } else {
            std::cerr << "Fork failed" << std::endl;
            return EXIT_FAILURE;
        }
    }

    // 最后1张卡使用两线程
    std::vector<std::thread> threads;
    int last_device_id = device_count - 1;
    // 创建线程执行最后一个设备的任务
    for (int i = 0; i < 2; ++i) {
        threads.emplace_back(gpuWorker, last_device_id);
    }
    // 等待所有线程完成
    for (auto& thread : threads) {
        thread.join();
    }

    // 等待所有子进程完成

    for (auto pid : pids) {
        int status;
        waitpid(pid, &status, 0);

        if (WEXITSTATUS(status) != EXIT_SUCCESS) {
            std::cerr << "Child process " << pid << " failed" << std::endl;
            return EXIT_FAILURE;
        }
    }

    std::cout << "All tests passed successfully!\n";
    return EXIT_SUCCESS;
}